#include"conv.h"
#include"compute.h"

void conv(At *in,Wt *weight,Ot *out,int fsize,int pad,int ch_in,int ch_out){
    int fsize_o=fsize+2*pad-K+1;
    for(int m=0;m<ch_out;m++)
    	for(int r=0;r<fsize_o;r++)
    		for(int c=0;c<fsize_o;c++){
    			Ot tmp=0;
    			for(int i=0;i<K;i++)
    				for(int j=0;j<K;j++)
    					for(int n=0;n<ch_in;n++)
    						if(r+i-pad>=0&&r+i-pad<fsize&&c+j-pad>=0&&c+j-pad<fsize)
    							tmp+=op(in[n*fsize*fsize+(r+i-pad)*fsize+(c+j-pad)],
    									weight[m*ch_in*K*K+n*K*K+i*K+j]);
    			out[m*fsize_o*fsize_o+r*fsize_o+c]=tmp;
    	}
}

void maxpool(Ot* in,Ot *out,int ch,int h,int w){
	for(int i=0;i<h/2;i++)
		for(int j=0;j<w/2;j++)
			for(int n=0;n<ch;n++){
				Ot max=-9999;
				for(int kx=0;kx<2;kx++)
					for(int ky=0;ky<2;ky++)
						if(*(in+n*h*w+(2*i+kx)*w+(2*j+ky))>max){
							max=*(in+n*h*w+(2*i+kx)*w+(2*j+ky));
						}
				*(out+n*h*w/4+i*w/2+j)=max;
			}
}

int main(){
	At in[N][12][12];
	At in_copy[N][12+2][12+2];       //after padding
	Ot out[M][12*12/4];
	Ot out_ref[M][12*12];
	Ot pool_out_ref[M][12*12/4];
	const Wt weight[M][N][K][K]={
#include"weight.h"
	};
	int pad=1;
	int fsize=12;
	//initialize
	for(int n=0;n<N;n++)
		for(int i=0;i<fsize;i++)
			for(int j=0;j<fsize;j++){
				in[n][i][j]=((n*fsize*fsize+i*fsize+j)%128-64);
			}
    //padding
	for(int n=0;n<N;n++)
		for(int i=0;i<fsize+2*pad;i++)
			for(int j=0;j<fsize+2*pad;j++)
				if(i>=pad&&j>=pad&&i<fsize+pad&&j<fsize+pad)
					in_copy[n][i][j]=in[n][i-pad][j-pad];
				else
					in_copy[n][i][j]=(At)0;
    //calculate
    conv((At*)in,(Wt*)weight,(Ot*)out_ref,fsize,pad,N,M);
    maxpool((Ot*)out_ref,(Ot*)pool_out_ref,M,fsize,fsize);
    //
    top((volatile At*) in_copy,(volatile At*) in_copy,(volatile Ot*)out,fsize+2*pad,N,M,1);

    int err=0;
    int correct=0;
    for(int m=0;m<M;m++)
    	for(int i=0;i<fsize*fsize/4;i++){
    		if(pool_out_ref[m][i]!=out[m][i])
    			err++;
    		else
    			correct++;
    		cout<<pool_out_ref[m][i]<<","<<out[m][i]<<endl;
    	}
    cout<<"error count is "<<err<<endl;
    cout<<"correct count is "<<correct<<endl;
    return 0;
}
